from fastapi import Body, Request
from fastapi.responses import StreamingResponse
from configs.model_config import LLM_MODELS, TEMPERATURE, HISTORY_LEN, MAX_TOKENS
from configs.kb_config import SCORE_THRESHOLD, VECTOR_SEARCH_TOP_K, GO_SERVICE_SECRET_KEY
from configs.other_config import INSTRUCTIONS, DEFAULT_CHARACTER, DEFAULT_NICKNAME, LOCATION_FILE_PATH, SECRET_KEY, ALGORITHM
from server.utils import BaseResponse
from configs.logging_config import configure_logging
from loguru import logger
from server.utils import create_chat_response_event, decode_token, decode_verify_token
from server.chat.utils import generate_request_id, log_content
from sse_starlette.sse import EventSourceResponse
from server.chat.utils import llm_chat
from server.db.repository.message_repository import get_message_by_user_id_mars_zhiliao, get_message_by_user_id_with_huohuo_chat, get_message_by_user_id_with_xingxing_chat
from typing import Optional
from jose import JWTError
import pytz
from server.chat.kb_chat_mars_zhiliao import knowledge_base_chat
from httpx import AsyncClient
from configs.kb_config import Go_SERVICE_URL_PREFIX
from datetime import datetime

knowledge_base_name = "Mars"
prompt_name = "knowledge_Mars"

configure_logging()
import json
with open(LOCATION_FILE_PATH, "r", encoding="utf-8") as f:
    locations = json.load(f)


async def mars_zhiliao_kb_chat(uid: str = Body(..., description="用户ID"),
                                query: str = Body(..., description="用户输入", examples=["恼羞成怒"]),
                                stream: bool = Body(False, description="流式输出"),
                                nickname: Optional[str] = Body(DEFAULT_NICKNAME, description="用户昵称"),
                                character: Optional[str] = Body(DEFAULT_CHARACTER, description="性格"),
                                request: Request = None
                                ):
    # 解码token
    token = request.headers.get("Authorization")
    if not token:
        return BaseResponse(code=401, msg="没有提供token")
    try:
        payload = decode_token(token, SECRET_KEY, ALGORITHM)
        token_uid = payload.get("sub")
        if not token_uid:
            return BaseResponse(code=401, msg="token无效")
        if token_uid != uid:
            return BaseResponse(code=401, msg="token无效")
    except JWTError as e:
        logger.error(f"Failed to decode token: {e}")
        return BaseResponse(code=401, msg="token无效", data={"error": str(e)})
    # 要求query 非空
    query = query.strip()
    if not query:
        return BaseResponse(code=400, msg="Query cannot be empty.")

    # 生成唯一ID
    request_id = generate_request_id()

    # 检查是否识别到了意图
    data_content = None
    try:
        instruction = await llm_chat(query, "Intention_Recognition_Expert", 0.1)
    except Exception as e:
        logger.error(f"Failed to recognize the instruction: {e}")
        instruction = "-1"
    city = ""
    district = ""
    if instruction in INSTRUCTIONS:
        answer = None
        request_id = "I_" + request_id
        if instruction == "5":
            try:
                info = await llm_chat(query, "Extract", 0.1)
                if info != "缺少信息":
                    info = info.split(",")
                    if len(info) == 1:
                        city = info[0]
                        district = ""
                    if len(info) == 2:
                        city = info[0]
                        district = info[1]
                if district == "None":
                    district = ""
            except Exception as e:
                logger.error(f"Failed to extract the information: {e}")
                request_id = "D_" + request_id
                # 获取uid的历史对话
                try:
                    if nickname == "火火":
                        histories = get_message_by_user_id_with_huohuo_chat(user_id=uid, history_len=HISTORY_LEN)
                    elif nickname == "星星":
                        histories = get_message_by_user_id_with_xingxing_chat(user_id=uid, history_len=HISTORY_LEN)
                except Exception as e:
                    logger.error(f"Failed to get chat history: {e}")
                    return BaseResponse(code=50003, msg="Failed to retrieve chat history.")
                response = await knowledge_base_chat(uid=uid,
                                                     query=query,
                                                     stream=stream,
                                                     history=histories,
                                                     knowledge_base_name=knowledge_base_name,
                                                     request_id=request_id,
                                                     score_threshold=SCORE_THRESHOLD,
                                                     top_k=VECTOR_SEARCH_TOP_K,
                                                     model_name=LLM_MODELS[0],
                                                     temperature=TEMPERATURE,
                                                     max_tokens=MAX_TOKENS,
                                                     request=request,
                                                     prompt_name=prompt_name,
                                                     nickname=nickname,
                                                     character=character)
                if stream:
                    responses = EventSourceResponse(response)
                else:
                    responses = StreamingResponse(response)
                return responses

        if stream:
            responses = EventSourceResponse(
                create_chat_response_event(response_text=answer, instruction=instruction, request_id=request_id,
                                           finish=True, data_content=data_content, city=city, district=district))
        else:
            responses = StreamingResponse(
                create_chat_response_event(response_text=answer, instruction=instruction, request_id=request_id,
                                           finish=True, data_content=data_content, city=city, district=district))
        log_content(uid, query, answer, instruction=instruction, request_id=request_id, city=city, district=district, chat_type="mars_zhiliao_kb_chat")
        return responses

    else:
        request_id = "D_" + request_id
        # 获取uid的历史对话
        try:
            if nickname == "火火":
                histories = get_message_by_user_id_with_huohuo_chat(user_id=uid, history_len=HISTORY_LEN)
            elif nickname == "星星":
                histories = get_message_by_user_id_with_xingxing_chat(user_id=uid, history_len=HISTORY_LEN)
        except Exception as e:
            logger.error(f"Failed to get chat history: {e}")
            return BaseResponse(code=50003, msg="Failed to retrieve chat history.")
        response = await knowledge_base_chat(uid=uid,
                                             query=query,
                                             stream=stream,
                                             history=histories,
                                             knowledge_base_name=knowledge_base_name,
                                             request_id=request_id,
                                             score_threshold=SCORE_THRESHOLD,
                                             top_k=VECTOR_SEARCH_TOP_K,
                                             model_name=LLM_MODELS[0],
                                             temperature=TEMPERATURE,
                                             max_tokens=MAX_TOKENS,
                                             request=request,
                                             prompt_name=prompt_name,
                                             nickname=nickname,
                                             character=character)
        if stream:
            responses = EventSourceResponse(response)
        else:
            responses = StreamingResponse(response)
        return responses



async def mars_zhiliao_kb_chat_v2(uid: str = Body(..., description="用户ID"),
                                query: str = Body(..., description="用户输入", examples=["恼羞成怒"]),
                                stream: bool = Body(False, description="流式输出"),
                                nickname: Optional[str] = Body(DEFAULT_NICKNAME, description="用户昵称"),
                                character: Optional[str] = Body(DEFAULT_CHARACTER, description="性格"),
                                occupation: Optional[str] = Body("", description="职业"),
                                request: Request = None
                                ):
    # 解码token
    try:
        token = decode_verify_token(request, GO_SERVICE_SECRET_KEY, ALGORITHM)
    except Exception as e:
        return BaseResponse(code=401, msg="Token is invalid", data={"error": str(e.detail.get('msg'))})
    
    if nickname == "星星":
        prompt_name = f"ZhiLiaoXingXing"
    elif nickname == "火火":
        prompt_name = f"ZhiLiaoHuoHuo"
    else:
        return BaseResponse(code=400, msg="昵称无效")
    # 要求query 非空
    query = query.strip()
    if not query:
        return BaseResponse(code=400, msg="Query cannot be empty.")

    # 生成唯一ID
    request_id = generate_request_id()

    # 检查是否识别到了意图
    data_content = None
    try:
        instruction = await llm_chat(query, "Intention_Recognition_Expert", 0.1)
    except Exception as e:
        logger.error(f"Failed to recognize the instruction: {e}")
        instruction = "-1"
    city = ""
    district = ""
    if instruction in INSTRUCTIONS:
        answer = None
        request_id = "I_" + request_id
        if instruction == "5":
            try:
                info = await llm_chat(query, "Extract", 0.1)
                if info != "缺少信息":
                    info = info.split(",")
                    if len(info) == 1:
                        city = info[0]
                        district = ""
                    if len(info) == 2:
                        city = info[0]
                        district = info[1]
                if district == "None":
                    district = ""
            except Exception as e:
                logger.error(f"Failed to extract the information: {e}")
                request_id = "D_" + request_id
                # 获取uid的历史对话
                try:
                    if nickname == "火火":
                        histories = get_message_by_user_id_with_huohuo_chat(user_id=uid, history_len=HISTORY_LEN)
                    elif nickname == "星星":
                        histories = get_message_by_user_id_with_xingxing_chat(user_id=uid, history_len=HISTORY_LEN)
                except Exception as e:
                    logger.error(f"Failed to get chat history: {e}")
                    return BaseResponse(code=50003, msg="Failed to retrieve chat history.")
                response = await knowledge_base_chat(uid=uid,
                                                     query=query,
                                                     stream=stream,
                                                     history=histories,
                                                     knowledge_base_name=knowledge_base_name,
                                                     request_id=request_id,
                                                     score_threshold=SCORE_THRESHOLD,
                                                     top_k=VECTOR_SEARCH_TOP_K,
                                                     model_name=LLM_MODELS[0],
                                                     temperature=TEMPERATURE,
                                                     max_tokens=MAX_TOKENS,
                                                     request=request,
                                                     prompt_name=prompt_name,
                                                     nickname=nickname,
                                                     character=character,
                                                     instruction=instruction)
                if stream:
                    responses = EventSourceResponse(response)
                else:
                    responses = StreamingResponse(response)
                return responses

        if stream:
            responses = EventSourceResponse(
                create_chat_response_event(response_text=answer, instruction=instruction, request_id=request_id,
                                           finish=True, data_content=data_content, city=city, district=district))
        else:
            responses = StreamingResponse(
                create_chat_response_event(response_text=answer, instruction=instruction, request_id=request_id,
                                           finish=True, data_content=data_content, city=city, district=district))
        log_content(uid, query, answer, instruction=instruction, request_id=request_id, city=city, district=district, chat_type="mars_zhiliao_kb_chat")
        return responses

    else:
        request_id = "D_" + request_id
        # 获取uid的历史对话
        try:
            if nickname == "火火":
                histories = get_message_by_user_id_with_huohuo_chat(user_id=uid, history_len=HISTORY_LEN)
            elif nickname == "星星":
                histories = get_message_by_user_id_with_xingxing_chat(user_id=uid, history_len=HISTORY_LEN)
        except Exception as e:
            logger.error(f"Failed to get chat history: {e}")
            return BaseResponse(code=50003, msg="Failed to retrieve chat history.")
        response = await knowledge_base_chat(uid=uid,
                                             query=query,
                                             stream=stream,
                                             history=histories,
                                             knowledge_base_name=knowledge_base_name,
                                             request_id=request_id,
                                             score_threshold=SCORE_THRESHOLD,
                                             top_k=VECTOR_SEARCH_TOP_K,
                                             model_name=LLM_MODELS[0],
                                             temperature=TEMPERATURE,
                                             max_tokens=MAX_TOKENS,
                                             request=request,
                                             prompt_name=prompt_name,
                                             nickname=nickname,
                                             character=character,
                                             instruction=instruction)
        if stream:
            responses = EventSourceResponse(response)
        else:
            responses = StreamingResponse(response)
        return responses


async def mars_zhiliao_kb_chat_v3(uid: str = Body(..., description="用户ID"),
                                query: str = Body(..., description="用户输入", examples=["恼羞成怒"]),
                                stream: bool = Body(False, description="流式输出"),
                                nickname: Optional[str] = Body(DEFAULT_NICKNAME, description="用户昵称"),
                                character: Optional[str] = Body(DEFAULT_CHARACTER, description="性格"),
                                occupation: Optional[str] = Body("", description="职业"),
                                request: Request = None
                                ):
    # 解码token
    try:
        token = decode_verify_token(request, GO_SERVICE_SECRET_KEY, ALGORITHM)
    except Exception as e:
        return BaseResponse(code=401, msg="Token is invalid", data={"error": str(e.detail.get('msg'))})
    
    # Go服务说1是火火，2是星星
    npc_id = -1
    if nickname == "星星":
        npc_id = 2
    if nickname == "火火":
        npc_id = 1
    # Go服务的获取职业接口
    try:
        url = f"{Go_SERVICE_URL_PREFIX}/admin/user/getnpcprofession"
        payload = {
            "UserId": int(uid),
            "NpcId": npc_id
        }
        async with AsyncClient() as client:
            response = await client.post(url, json=payload, timeout=30)
            # 过期时间
            profession = response.json().get("Profession")
            if profession != None:
                expiration_time = profession.get("ExpireTime") # int时间戳
                # 判断是否过期
                if expiration_time:
                    # 创建北京时区对象
                    tz_beijing = pytz.timezone("Asia/Shanghai")
                    # 获取当前时间戳
                    now = datetime.now(tz_beijing)
                    if expiration_time < now.timestamp():
                        is_expired = True
                    else:
                        is_expired = False
                else:
                    logger.warning(f"ExpireTime 字段为空")
                    is_expired = True
            else:
                is_expired = True
    except Exception as e:
        logger.error(f"Failed to get occupation: {e}")
        return BaseResponse(code=50006, msg="职业获取失败")

    # 如果职业不为空，则使用职业版prompt
    if not is_expired and occupation != "":
        if nickname == "星星":
            prompt_name = f"ZhiLiaoXingXingOccupation"
        elif nickname == "火火":
            prompt_name = f"ZhiLiaoHuoHuoOccupation"
        else:
            return BaseResponse(code=400, msg="有职业，但是昵称无效")
    else:
        if nickname == "星星":
            prompt_name = f"ZhiLiaoXingXing"
        elif nickname == "火火":
            prompt_name = f"ZhiLiaoHuoHuo"
        else:
            return BaseResponse(code=400, msg="没有职业，但是昵称无效")
    # 要求query 非空
    query = query.strip()
    if not query:
        return BaseResponse(code=400, msg="Query cannot be empty.")

    # 生成唯一ID
    request_id = generate_request_id()

    # 检查是否识别到了意图
    data_content = None
    try:
        instruction = await llm_chat(query, "Intention_Recognition_Expert", 0.1)
    except Exception as e:
        logger.error(f"Failed to recognize the instruction: {e}")
        instruction = "-1"
    city = ""
    district = ""
    if instruction in INSTRUCTIONS:
        answer = None
        request_id = "I_" + request_id
        if instruction == "5":
            try:
                info = await llm_chat(query, "Extract", 0.1)
                if info != "缺少信息":
                    info = info.split(",")
                    if len(info) == 1:
                        city = info[0]
                        district = ""
                    if len(info) == 2:
                        city = info[0]
                        district = info[1]
                if district == "None":
                    district = ""
            except Exception as e:
                logger.error(f"Failed to extract the information: {e}")
                request_id = "D_" + request_id
                # 获取uid的历史对话
                try:
                    if nickname == "火火":
                        histories = get_message_by_user_id_with_huohuo_chat(user_id=uid, history_len=HISTORY_LEN)
                    elif nickname == "星星":
                        histories = get_message_by_user_id_with_xingxing_chat(user_id=uid, history_len=HISTORY_LEN)
                except Exception as e:
                    logger.error(f"Failed to get chat history: {e}")
                    return BaseResponse(code=50003, msg="Failed to retrieve chat history.")
                
                response = await knowledge_base_chat(uid=uid,
                                                     query=query,
                                                     stream=stream,
                                                     history=histories,
                                                     knowledge_base_name=knowledge_base_name,
                                                     request_id=request_id,
                                                     score_threshold=SCORE_THRESHOLD,
                                                     top_k=VECTOR_SEARCH_TOP_K,
                                                     model_name=LLM_MODELS[0],
                                                     temperature=TEMPERATURE,
                                                     max_tokens=MAX_TOKENS,
                                                     request=request,
                                                     prompt_name=prompt_name,
                                                     nickname=nickname,
                                                     character=character,
                                                     instruction=instruction)
                if stream:
                    responses = EventSourceResponse(response)
                else:
                    responses = StreamingResponse(response)
                return responses

        if stream:
            responses = EventSourceResponse(
                create_chat_response_event(response_text=answer, instruction=instruction, request_id=request_id,
                                           finish=True, data_content=data_content, city=city, district=district))
        else:
            responses = StreamingResponse(
                create_chat_response_event(response_text=answer, instruction=instruction, request_id=request_id,
                                           finish=True, data_content=data_content, city=city, district=district))
        log_content(uid, query, answer, instruction=instruction, request_id=request_id, city=city, district=district, chat_type="mars_zhiliao_kb_chat")
        return responses

    else:
        # 星星是律师，火火是情感专家, 如果nickname不是星星，但是识别到的指令是1000，那么指令初始化为-1
        if nickname != "星星" and instruction == "1000":
            instruction = "-1"
        # 如果nickname不是火火，但是识别到的指令是1001，那么指令初始化为-1
        if nickname != "火火" and instruction == "1001":
            instruction = "-1"
        request_id = "D_" + request_id
        # 获取uid的历史对话
        try:
            if nickname == "火火":
                histories = get_message_by_user_id_with_huohuo_chat(user_id=uid, history_len=HISTORY_LEN)
            elif nickname == "星星":
                histories = get_message_by_user_id_with_xingxing_chat(user_id=uid, history_len=HISTORY_LEN)
        except Exception as e:
            logger.error(f"Failed to get chat history: {e}")
            return BaseResponse(code=50003, msg="Failed to retrieve chat history.")
        response = await knowledge_base_chat(uid=uid,
                                             query=query,
                                             stream=stream,
                                             history=histories,
                                             knowledge_base_name=knowledge_base_name,
                                             request_id=request_id,
                                             score_threshold=SCORE_THRESHOLD,
                                             top_k=VECTOR_SEARCH_TOP_K,
                                             model_name=LLM_MODELS[0],
                                             temperature=TEMPERATURE,
                                             max_tokens=MAX_TOKENS,
                                             request=request,
                                             prompt_name=prompt_name,
                                             nickname=nickname,
                                             character=character,
                                             instruction=instruction)
        if stream:
            responses = EventSourceResponse(response)
        else:
            responses = StreamingResponse(response)
        return responses